import imgaug.augmenters as iaa
from aug_tools.transforms import *  # 这里的transforms是上面的transforms.py文件，如果修改了文件名，这里对应修改即可


# imgaug是一个用于机器学习实验中图像增强的python库

class RotateAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([  # 定义变换序列, 可根据需要自行增减或修改参数
            # iaa.Dropout([0.0, 0.01]),  # 随机去掉一些像素点，即把这些像素点变成0
            #  iaa.Sharpen((0.0, 0.2)),  # 锐化处理
             iaa.Affine(rotate=(-15, 15)), #旋转
            # 仿射变换， rotate by -45 to 45 degrees (affects segmaps)
            #iaa.AddToBrightness((-30, 30))  # 改变亮度
            # iaa.AddToHue((-20, 20)),  # 色调随机
            # iaa.Sometimes(0.5, iaa.Fliplr(1)),  0.5概率翻转
            # iaa.Fliplr(p=1),  # 翻转图片，水平翻转图像（左右）
        ], random_order=True)


AUGMENTATION_TRANSFORMS_Rotate = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    RotateAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])

class FliplrAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([  # 定义变换序列, 可根据需要自行增减或修改参数

             iaa.Fliplr(p=1)

        ], random_order=True)


AUGMENTATION_TRANSFORMS_Fliplr = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    FliplrAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])

class FlipudAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([  # 定义变换序列, 可根据需要自行增减或修改参数

             iaa.Flipud(p=1)

        ], random_order=True)


AUGMENTATION_TRANSFORMS_Flipud = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    FlipudAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])

class BlurAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([  # 定义变换序列, 可根据需要自行增减或修改参数
             iaa.GaussianBlur(sigma=(0.25, 1.75)),
             # iaa.AverageBlur(k=(1, 5)),
             # iaa.MedianBlur(k=(1, 5))
        ], random_order=True)


AUGMENTATION_TRANSFORMS_Blur = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    BlurAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])

class GaussianNoiseAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([  # 定义变换序列, 可根据需要自行增减或修改参数

             iaa.AdditiveGaussianNoise(10,20)

        ], random_order=True)


AUGMENTATION_TRANSFORMS_GaussianNoise = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    GaussianNoiseAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])

class HueAndSaturationAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([  # 定义变换序列, 可根据需要自行增减或修改参数

             iaa.AddToHueAndSaturation(value_hue=-20,value_saturation=20)

        ], random_order=True)


AUGMENTATION_TRANSFORMS_HueAndSaturation = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    HueAndSaturationAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])

class BrightnessAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([  # 定义变换序列, 可根据需要自行增减或修改参数

            iaa.AddToBrightness((-30, 30))  # 改变亮度

        ], random_order=True)


AUGMENTATION_TRANSFORMS_Brightness = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    BrightnessAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])


class ContrastAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([  # 定义变换序列, 可根据需要自行增减或修改参数

            iaa.contrast.LinearContrast(alpha=(0.6, 1.4))

        ], random_order=True)


AUGMENTATION_TRANSFORMS_Contrast = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    ContrastAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])

class RandomAug(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.SomeOf(np.random.randint(1,10),[  # 定义变换序列, 可根据需要自行增减或修改参数

            # iaa.Noop(),
            iaa.Affine(rotate=(-15, 15)),
            iaa.Fliplr(p=1),
            iaa.Flipud(p=1),
            iaa.GaussianBlur(sigma=(0.25, 1.75)),
            iaa.AverageBlur(k=(1, 5)),
            iaa.MedianBlur(k=(1, 5)),
            iaa.AdditiveGaussianNoise(10, 20),
            iaa.AddToHueAndSaturation(value_hue=-20, value_saturation=20),
            iaa.AddToBrightness((-30, 30)),
            iaa.contrast.LinearContrast(alpha=(0.6, 1.4)),


        ],random_order=True)


AUGMENTATION_TRANSFORMS_Random = transforms.Compose([
    AbsoluteLabels(),  # 绝对标签
    RandomAug(),  # 一些基本的数据增强
    RelativeLabels(),  # 相对标签
])
